import os
import json
import configparser
import asyncio
import websockets

from typing import Any, Optional, Union
from sparkai.llm.llm import ChatSparkLLM, AsyncChunkPrintHandler
from sparkai.core.outputs import (
    ChatGenerationChunk,
    GenerationChunk,
)

config = configparser.ConfigParser()
config.read('xunfei.ini')

SPARKAI_URL = config['DEFAULT'].get('SPARKAI_URL')
SPARKAI_APP_ID = config['DEFAULT'].get('SPARKAI_APP_ID')
SPARKAI_API_SECRET = config['DEFAULT'].get('SPARKAI_API_SECRET')
SPARKAI_API_KEY = config['DEFAULT'].get('SPARKAI_API_KEY')
SPARKAI_DOMAIN = config['DEFAULT'].get('SPARKAI_DOMAIN')



class MyHandler(AsyncChunkPrintHandler):
    def __init__(self, color: Optional[str] = None, websocket = None) -> None:
        """Initialize callback handler."""
        super().__init__()
        self.color = color
        self.websocket = websocket
    async def on_llm_new_token(
            self,
            token: str,
            *,
            chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
            **kwargs: Any,
    ) -> None:

        run_id = str(kwargs.get("run_id", ""))
        # 是否结束 final
        final = kwargs.get("final", False)
        await self.websocket.send(json.dumps({"token": token, "run_id": run_id, 'final': final}))

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=True,
    max_tokens=1024,
)

async def send_to_spark(websocket, messages):
    handler = MyHandler(websocket=websocket)
    a = spark.astream(messages, config={"callbacks": [handler]})
    async for send_msg in a:
        print("===========================================================================")
        print(send_msg.content)
        print("===========================================================================")

class MyWebSocketServer:
    def __init__(self, host='0.0.0.0', port=8765):
        self.host = host
        self.port = port
        self.server = None

    async def handle_client(self, websocket):
        print("已连接")
        print(f"信息： {websocket.remote_address}")
        print("等待发送消息...")
        messageMap = {}
        try:
            while True:
                clientMsg = await websocket.recv()
                print(f"收到消息.")
                clientMsg = json.loads(clientMsg)
                print(clientMsg)
                mkey = clientMsg['session_code']
                # 初次通信 把历史消息列表存入 messageMap
                if clientMsg['type'] == 'init':
                    messageMap[mkey] = clientMsg['data']
                else:
                    messageMap[mkey].append({
                        "role": "user",
                        "content": clientMsg['data']
                    })
                # print(messageMap)
                # 消息不为空 且 最后一个消息不是AI回答的
                if len(messageMap[mkey]) > 0 and messageMap[mkey][-1]['role'] == 'user':
                    # 发送给LLM
                    await send_to_spark(websocket, messageMap[mkey])
                # 空消息 根据角色属性 发起对话
                if len(messageMap[mkey]) == 0:
                    contentStart = f"你将扮演一个 {clientMsg['role_name']} 角色, 你的性别是: {clientMsg['role_sex']}, 你的年龄是: {clientMsg['role_age']}，你的性格是: {clientMsg['role_character']}， 你现在正在和用户对话，请回答用户的问题。"
                    messageMap[mkey].append({
                        "role": "user",
                        "content": contentStart
                    })
                    print("contentStart: " + contentStart)
                    await send_to_spark(websocket, messageMap[mkey])
        except websockets.ConnectionClosedOK:
            print("WebSocket连接已正常关闭")
        except websockets.ConnectionClosedError as e:
            print(f"WebSocket连接已关闭：{e}")
        except Exception as e:
            print(f"发生错误：{e}")
        finally:
            print("客户端已断开连接")
    async def start(self):
        async with websockets.serve(self.handle_client, self.host, self.port):
            print(f"WebSocket 服务器已启动，监听地址：{self.host}:{self.port}")
            await asyncio.Future()
    async def stop(self):
        await self.server.wait_closed()
        print("WebSocket 服务器已停止")


if __name__ == '__main__':
    websocket_server = MyWebSocketServer('0.0.0.0', 9999)
    asyncio.run(websocket_server.start())
